package org.nativejpa.transaction;

import org.apache.commons.logging.LogFactory;
import org.apache.commons.logging.Log;

import javax.transaction.Transaction;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.ArrayList;
import java.util.Map;
import java.util.HashMap;

/**
 * @author <a href="mailto:jfox.young@gmail.com">Young Yang</a>
 * @create 2009-3-17 17:15:05
 */
public class TransactionConnectionManager {
    private static Log log = LogFactory.getLog(TransactionConnectionManager.class);

    /**
     * ThreadLocal 保存 Thread => TransactionConnectionPair
     */
    private static ThreadLocal<TxConnectionHolder> tx2Connections = new ThreadLocal<TxConnectionHolder>() {
        @Override
        protected TxConnectionHolder initialValue() {
            return new TxConnectionHolder();
        }
    };

    public static void addConnection2Tx(String unitName, Connection connection, Transaction transaction) {
        log.debug("Attatch connection to transaction, connection is: " + connection);
        TxConnectionHolder txConnectionHolder =  tx2Connections.get();

        if(txConnectionHolder == null) {
            log.warn("Attatch connection to transaction failed, no transaction!");
        }
        else {
            txConnectionHolder.addConnection(unitName, connection, transaction);
        }
    }

    /**
     * reuse connection is 
     * @param unitName
     * @return
     */
    public static Connection reuseTransactionConnection(String unitName, Transaction transaction){
        TxConnectionHolder txConnectionHolder =  tx2Connections.get();
        if(txConnectionHolder != null) {
            return txConnectionHolder.reuseConnectionByUnit(unitName, transaction);
        }
        return null;
    }

    public static void releaseTxConnections(Transaction transaction){
        log.debug("Release Transaction attatched connections.");
        TxConnectionHolder txConnectionHolder = tx2Connections.get();
        if(txConnectionHolder == null){
            log.warn("Release attatched connections failed, current transaction is null.");
        }
        else {
            txConnectionHolder.releaseCurrentTxConnections(transaction);
            // 如果没有剩余的事务环境，清除 ThreadLocal 变量
            if(txConnectionHolder.isEmpty()) {
                log.debug("Remove thread local transaction binding.");
                tx2Connections.remove();
            }
        }
    }

    /**
     * 保存和 Thread 关联的 Transaction 以及 Connection
     */
    static class TxConnectionHolder {

        private Map<Transaction, Map<String,UnitConnectionHolder>> tx2ConnectionsMap = new HashMap<Transaction, Map<String, UnitConnectionHolder>>(3);

/*
        public void addTransaction(Transaction transaction){
            transactionChain.add(transaction);
        }
*/

        public void addConnection(String unitName, Connection connection, Transaction transaction){
            if(!tx2ConnectionsMap.containsKey(transaction)) {
                UnitConnectionHolder uch = new UnitConnectionHolder(unitName);
                Map<String, UnitConnectionHolder> unit2Connections = new HashMap<String, UnitConnectionHolder>(3);
                unit2Connections.put(unitName,uch);
                tx2ConnectionsMap.put(transaction, unit2Connections);
            }            
            UnitConnectionHolder uch = tx2ConnectionsMap.get(transaction).get(unitName);
            uch.addConection(connection);
        }

        public Connection reuseConnectionByUnit(String unitName, Transaction transaction) {
             Map<String, UnitConnectionHolder> unit2Connections = tx2ConnectionsMap.get(transaction);
            if(unit2Connections != null && unit2Connections.containsKey(unitName)) {
                return unit2Connections.get(unitName).getConnection();
            }
            return null;
        }

        /**
         * 释放当前活动事务关联所有 connection
         */
        public void releaseCurrentTxConnections(Transaction transaction){
            Map<String, UnitConnectionHolder> unit2Connections = tx2ConnectionsMap.remove(transaction);
            if(unit2Connections != null && !unit2Connections.isEmpty()) {
                for(UnitConnectionHolder uch : unit2Connections.values()) {
                    for(Connection conn : uch.getConnections()){
                        try {
                            conn.close();
                        }
                        catch (SQLException e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
        }

        public boolean isEmpty(){
            return tx2ConnectionsMap.isEmpty();
        }
    }

    static class UnitConnectionHolder {
        private String unitName;
        // normarly, only one connection of one entitymanager in one transaction
        private List<Connection> connections = new ArrayList<Connection>(2);

        UnitConnectionHolder(String unitName) {
            this.unitName = unitName;
        }

        void addConection(Connection connection) {
            if(!connections.contains(connection)) {
                connections.add(connection);
            }
        }

        Connection getConnection(){
            if(!connections.isEmpty()) {
                return connections.get(connections.size()-1);
            }
            return null;
        }

        Connection[] getConnections(){
            return connections.toArray(new Connection[connections.size()]);
        }

    }

}
